# coding: utf-8
import os.path
import time

import matplotlib.pyplot as plt
import torch
from communication import *
from commpy.channelcoding import get_ldpc_code_params
from utils import *

# TEMP STORAGE
QAM_BER = []
LDPC_BER = []

# def compute_ber(origin_bits, target_bits):
#     return np.mean(np.abs(origin_bits - target_bits))

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++#
#                                   Simulate Communication                                           #
#           image -> bits -> LDPC -> QAM -> Channel -> QAM -> LDPC -> bits -> rec_image              #
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++#


class FastLdpcQamChannel:
    def __init__(self,
                 block_size=1440*1e4,
                 ldpc_param_path="./channel_codec/1440.720.txt",
                 qam_order=16,
                 channel_type="awgn",
                 snr=5,
                 device=None
                 ):
        self.device = device if device is not None else "cuda" if torch.cuda.is_available() else "cpu"
        self.param = get_ldpc_code_params(ldpc_param_path, compute_matrix=True)
        self.qam_order = qam_order
        self.qam = QAMModem(qam_order, self.device)
        self.block_size = int(block_size)
        self.channel_type = channel_type
        # function
        # self.channel = {
        #     "awgn": awgn_torch,
        #     "rayleigh": rayleigh_torch,
        #     "rician": rician_torch
        # }[self.channel_type]
        # class
        self.channel = Channel(self.channel_type, device)
        self.snr = snr

    def img2bits(self, image):
        return image2bits(image)

    def bits2img(self, bits, shape):
        return bits2image(bits, shape)

    def forward(self, bits):
        bits = torch.tensor(bits, device=self.device)
        # 1 ldpc encode
        ldpc_encode_bits = triang_ldpc_systematic_encode_torch(bits, self.param)
        first_dimension_length = ldpc_encode_bits.shape[0]
        flatten_bits = ldpc_encode_bits.reshape(-1)
        # 2 TODO qam modulate
        # modulated_bits = flatten_bits
        modulated_bits = self.qam.modulate(flatten_bits)

        # 3 channel
        bits_with_noise = self.channel(modulated_bits, self.snr)
        # TODO implemented by commpy
        # numpy_bits = modulated_bits.cpu().detach().numpy()
        # bits_with_noise = commpy.awgn(numpy_bits, self.snr)
        # bits_with_noise = torch.tensor(bits_with_noise, device=self.device)
        # 4 TODO qam demodulate
        # demodulated_bits = bits_with_noise
        demodulated_bits = self.qam.demodulate(bits_with_noise, "hard")

        # print(f"QAM BER: {(flatten_bits != demodulated_bits).sum()/len(demodulated_bits)}")
        # QAM_BER.append(((flatten_bits != demodulated_bits).sum()/len(demodulated_bits)).cpu().item())
        # 5 ldpc decode
        ldpc_decode_bits = demodulated_bits.reshape(first_dimension_length, -1)
        ldpc_decode_bits = 1 - 2 * ldpc_decode_bits
        ldpc_decode_bits = torch.tensor(ldpc_decode_bits.detach().cpu().numpy().reshape(-1, order='F').astype(float), device=self.device)
        ldpc_decoded_bits = torch_ldpc_bp_decode(ldpc_decode_bits,
                                                 self.param,
                                                 'MSA',
                                                 20)[0][:720]
        # print(f"LDPC BER: {(bits.detach().cpu().numpy() != ldpc_decoded_bits.detach().cpu().numpy().reshape(-1, order='F')[:len(bits)]).sum()/len(bits)}")
        # LDPC_BER.append(((bits.detach().cpu().numpy() != ldpc_decoded_bits.detach().cpu().numpy().reshape(-1, order='F')[:len(bits)]).sum()/len(bits)))
        return ldpc_decoded_bits.detach().cpu().numpy().reshape(-1, order='F')[:len(bits)]

    def __call__(self, image):
        t1 = time.time()
        bits, shape = self.img2bits(image)
        # block
        bits_length = len(bits)
        receive_data = []
        if bits_length > self.block_size:
            num = bits_length // self.block_size
            data_list = [bits[i:i+self.block_size] for i in range(0, bits_length, self.block_size)]
            if bits_length % self.block_size != 0:
                data_list.append(bits[num*self.block_size:])
            for data in data_list:
                receive_data.extend(
                    self.forward(data)
                )
        else:
            receive_data.extend(
                self.forward(bits)
            )
        rec_image = self.bits2img(receive_data, shape)
        t2 = time.time()
        print(f"Torch LDPC+{self.qam_order}-QAM+{self.channel_type} consuming: {t2-t1} seconds.")
        return rec_image


if __name__ == '__main__':
    from PIL import Image
    from snr_test import plot_snr_test
    # seed = 2025
    # torch.random.manual_seed(seed)
    # np.random.seed(seed)
    start = -5
    end = 26
    step = 2
    size = "512x768"
    # size = "1339x2048"
    qam_order = 64
    # pass channel
    for snr in range(start, end, step):
        # image_path = r"C:\Users\13541\Desktop\alejandro-escamilla-6_dec.png"   # HxW (1339, 2048)
        image_path = r"D:\data\cifar-10-batches-py\temp_test\33\kodim23_dec.png"  # HxW (512, 768)
        img = Image.open(image_path)
        model = FastLdpcQamChannel(ldpc_param_path=r"F:\research-code\2025\H256+LDPC\channel_codec\1440.720.txt",
                                    qam_order=qam_order, channel_type="awgn",
                                    snr=snr)
        rec_image = model(img)
        img_save_path = r"F:\PCL\month\2025\25-08\experiments\H.265\%s\%s_%dqam_snr%d.png" \
              % (size, os.path.basename(image_path).strip("_dec.png"), qam_order, snr)
        print(img_save_path)
        rec_image.save(img_save_path)

    # origin_image_path = r"D:\video-communication-dataset\图像编解码测试\CLIC\professional_valid\alejandro-escamilla-6.png"
    origin_image_path = r"D:\video-communication-dataset\图像编解码测试\kodak\archive\kodim23.png"

    snr_image_folder = r"F:\PCL\month\2025\25-08\experiments\H.265\%s" % size
    plot_snr_test(origin_image_path, snr_image_folder, start, end, step, qam_order,
                  save_path=os.path.join(snr_image_folder, "test_%s.png" % size))
    print("Successfully!!!")
    # plt.figure()
    # plt.title("1/2LDPC")
    # plt.xlabel("SNR(dB)")
    # plt.ylabel("BER")
    # plt.plot(list(range(start, end, step)), LDPC_BER, "bo-")
    # plt.show()
    pass

